'''
Created on Jul 7, 2010

@author: oabalbin
'''
import sys
import glob
import os
import numpy as np
import scipy as sp
from scipy import stats
from collections import deque, defaultdict

from RNAseq.common.classes import myArray


def open_files(folderpath,filename):
    '''
    Open the file with the expression for the kinases and other genes.
    '''
    myfiles=[]
    for infile in glob.glob( os.path.join(folderpath, filename) ):
        myfiles.append(infile)
    
    return myfiles

def list_of_names(inputfile):
        
        genlist=deque()
        for line in inputfile:        
            line = line.strip('\n')
            fields = line.split('\t')
            #print fields 
            # To skip headers, star reading samples in column 7 of the file
            if fields[0][0] == '#': 
                continue                             
            if fields[0]=='':
                continue
            
            genlist.append(fields[0])
        
        return genlist

def list_of_samples(inputfile_name):
        
        inputfile = open(inputfile_name)        
        samplelist, normalslist=deque(), deque()
        for line in inputfile:
            if line.startswith('#'):
                fields = line.strip('#\n').split('\t')
                if fields[0] == 'Samples':
                    Normals=False
                    continue
                elif fields[0] == 'Normals':
                    Normals=True
                    continue
            
            if not Normals:
                fields = line.strip('\n').split('\t')
                samplelist.append(fields[0])
            else:
                fields = line.strip('\n').split('\t')
                normalslist.append(fields[0])
        
        inputfile.close()
        
        return samplelist, normalslist


def dict_of_names(inputfile):
        
        genlist=defaultdict()
        i=0
        for line in inputfile:        
            line = line.strip('\n')
            fields = line.split('\t')
            #print fields 
            # To skip headers, star reading samples in column 7 of the file
            if fields[0][0] == '#': 
                continue                             
            if fields[0]=='':
                continue
            
            genlist[fields[0]] = i
            i+=1 
        
        return genlist


def create_myArray(myfile, median_threshold, normal_samples):
    '''
    It creates a sample object and calculates a z-score for the expression level of each gene
    in the sample. Z-score (expvalues - genemedian)/genesd
    Returns a geneArray and also a samplelist with the expression for each sample
    The input matrix is already lo2 normalized    
    '''
    gene_dictionary, sample_dictionary, gene_interval_dict = defaultdict(), defaultdict(), defaultdict()
    celllines =  deque()
    normal_samples_list=[]
    
    first=True

    for line in myfile:
        if line is None:
            continue 
        if line.startswith('#'):
            continue
        # Handeling headers ....
       
        if line.startswith('gene'):
            celllines = line.strip('\n').split('\t')[2:]
            for j,sp in enumerate(celllines):
                sample_dictionary[j] = sp
                print sp,j
                if sp in normal_samples:
                    normal_samples_list.append(j)
                    
            k=0
            continue
        
        # Reading the file.
        fields = line.strip('\n').split('\t') 
        gene_dictionary[k] = fields[0]
        gene_interval_dict[k] = fields[1]
        
        exp=np.zeros(len(fields[2:]))
        expmasked = np.ma.make_mask_none((len(fields[2:]),))        
        exp = np.array(map(float,fields[2:]))
        
        # shift the values by one. In order to avoid the zero coverage values which will be inf when log is applied.    
        expvalues = np.array([exp]) + 1.0
        expvalues = np.ma.array(expvalues,mask=[expmasked])
        
        if first==True:
            geneRawValues = np.array(expvalues)
            first=False
        else:
            geneRawValues = np.concatenate((geneRawValues,expvalues),axis=0)
        
        k+=1
    
    # Pre-process the matrix data in order to filter out low expressing genes.
    # Calculate the median expression value for each gene.
    gene2keep = filter_low_expressing_genes(geneRawValues,median_threshold)
    geneRawValues = geneRawValues[gene2keep,:]
    gene_dictionary = rebuild_gene_dictionary(gene_dictionary,gene2keep)
    gene_interval_dict = rebuild_gene_dictionary(gene_interval_dict,gene2keep)
    
    # Create an GeneArray object
    thisGeneArray = myArray(geneRawValues, gene_dictionary, sample_dictionary, normal_samples_list, gene_interval_dict)

    return thisGeneArray


def filter_low_expressing_genes(geneRawValues,threshold):
    """
    It filters out genes with very low expression values. 
    genes which max expression value  < threshold for all samples are removed.          
    """
    print geneRawValues.shape[0]
    gene_indicators = np.arange(geneRawValues.shape[0])
    exp_cutoff = np.ma.median(geneRawValues, axis=1)
    #median_cutoff = sp.stats.scoreatpercentile(gene_medians,threshold) 
    #exp_cutoff = np.ma.max(geneRawValues, axis=1)
    gene2keep = gene_indicators[exp_cutoff > threshold]
    print len(gene2keep)
    return gene2keep

def rebuild_gene_dictionary(genedict,genes2keep):
    """
    """
    new_gene_dictionary=defaultdict()
    for i in range(len(genes2keep)):
        new_gene_dictionary[i] = genedict[genes2keep[i]]
    
    return new_gene_dictionary


def print_prot_expression_matrix_stdfiltered(thisGeneArray, outfile_matrix, lstd):
    """
    """
    list_of_samples = thisGeneArray.get_list_of_samples()
    #print list_of_samples
    list_of_genes = thisGeneArray.get_list_of_genes()
    outfile_matrix.write("protid"+'\t'+"protLen"+'\t'+",".join(list_of_samples).replace(',','\t')+'\n')
    
    gene_std = thisGeneArray.sdv_array()
    gene_med = thisGeneArray.median_array()
    std_perc = sp.stats.scoreatpercentile(gene_std,lstd)
    med_perc = sp.stats.scoreatpercentile(gene_med,lstd)
    
    for i in range(thisGeneArray.expVal.shape[0]):
        
        if (gene_std[i] < std_perc) or (gene_med[i] < med_perc):
            continue
        eline = list(thisGeneArray.expVal[i,:])
        
        # especially for proteomics experimens        
        nzeros = eline.count(0.0)
        if nzeros == len(eline):
            continue
        #eline = list(thisGeneArray.expVal_foldchange_mean[i,:])
        length = thisGeneArray.genInterval[i] 
        outfile_matrix.write(list_of_genes[i]+'\t'+str(length)+'\t'+",".join(map(str,eline)).replace(',','\t')+'\n')

def print_prot_expression_matrix_stdfiltered_nsaf(thisGeneArray, outfile_matrix, lstd):
    """
    """
    list_of_samples = thisGeneArray.get_list_of_samples()
    #print list_of_samples
    list_of_genes = thisGeneArray.get_list_of_genes()
    outfile_matrix.write("protid"+'\t'+"protLen"+'\t'+",".join(list_of_samples).replace(',','\t')+'\n')
    
    gene_std = thisGeneArray.sdv_array()
    gene_med = thisGeneArray.median_array()
    std_perc = sp.stats.scoreatpercentile(gene_std,lstd)
    med_perc = sp.stats.scoreatpercentile(gene_med,lstd)
    
    for i in range(thisGeneArray.expVal_nsaf.shape[0]):
        
        if (gene_std[i] < std_perc) or (gene_med[i] < med_perc):
            continue
        eline = list(thisGeneArray.expVal_nsaf[i,:])

        length = thisGeneArray.genInterval[i] 
        outfile_matrix.write(list_of_genes[i]+'\t'+str(length)+'\t'+",".join(map(str,eline)).replace(',','\t')+'\n')
